﻿#ifndef __JS_COLOR_MODEL_H__
#define __JS_COLOR_MODEL_H__

#include "models/ColorModel.h"
using namespace darknet;

class JsColorModel : public ColorModel
{
public:
    static JsColorModel *GetInstance()
    {
        static JsColorModel s_instance;
        return &s_instance;
    }

    JsColorModel(const std::string& model_name = "colors")
    {
        m_model_name = model_name;

#ifdef __EMSCRIPTEN__
        std::string root_directory = "/assets/colors";
#else
        std::string root_directory = "D:/workspace/cpp/darknet/web/asset/assets/colors";
#endif
        std::string model_directory   = root_directory + "/models";
        std::string model_filepath    = model_directory + "/" + GetDefaultModelFileName();
        std::string dataset_directory = root_directory + "/datasets";
        std::string dataset_filepath  = dataset_directory + "/" + GetModelName() + ".dataset";

        SetRootDirectory(root_directory);
        SetModelDirectory(model_directory);
        SetModelPath(model_filepath);
        SetDataSetDirectory(dataset_directory);
        SetDataSetPath(dataset_filepath);
        SetDataSetColorsPath(dataset_directory + "/" + GetModelName() + ".dataset.colors");
        
#ifdef __EMSCRIPTEN__
        LoadDataSet();
        // std::cout << "load default dataset: " << GetDataSetPath() << " size = " << GetDataSet()->size() << std::endl;
#else
        std::filesystem::create_directories(root_directory);
        std::filesystem::create_directories(model_directory);
        std::filesystem::create_directories(dataset_directory);
#endif

        if (std::filesystem::exists(model_filepath) && LoadModelFromFile(model_filepath))
        {
            std::cout << "load prebuilt model from " << model_filepath << std::endl;
        }
        else
        {
            LoadNet();
        }
        SetModelDirectory(model_directory);
    }

    int Predict(unsigned char r, unsigned char g, unsigned char b)
    {
        return ColorModel::Predict(r, g, b);
    }

    int *GetColorRect(std::shared_ptr<DataImage> img, int color_classify_index, float scale = 0.2)
    {
        img = img->Scale(scale, scale);

        std::vector<int> xs;
        std::vector<int> ys;

        auto img_width    = img->GetWidth();
        auto img_height   = img->GetHeight();
        auto img_channels = img->GetChannel();
        auto img_datas    = img->GetData();

        for (int y = 0; y < img_height; y++)
        {
            for (int x = 0; x < img_width; x++)
            {
                auto r     = img->GetPixel(x, y, 0) * 255;
                auto g     = img->GetPixel(x, y, 1) * 255;
                auto b     = img->GetPixel(x, y, 2) * 255;
                auto index = Predict(r, g, b);
                if (color_classify_index == index)
                {
                    xs.push_back(x);
                    ys.push_back(y);
                }
            }
        }

        int size = xs.size();
        if (size == 0) return nullptr;

        auto avg_width  = 0;
        auto avg_height = 0;
        for (int i = 0; i < size; i++)
        {
            auto max_width  = 0;
            auto max_height = 0;
            for (int j = i + 1; j < size; j++)
            {
                max_width  = std::max(max_width, std::abs(xs[i] - xs[j]));
                max_height = std::max(max_height, std::abs(ys[i] - ys[j]));
            }
            avg_width += max_width;
            avg_height += max_height;
        }
        avg_width  = avg_width / size * 2;
        avg_height = avg_height / size * 2;

        auto avg_x      = 0;
        auto avg_y      = 0;
        auto min_x      = img_width;
        auto min_y      = img_height;
        auto max_x      = 0;
        auto max_y      = 0;
        auto valid_size = 0;
        for (int i = 0; i < size; i++)
        {
            auto x     = xs[i];
            auto y     = ys[i];
            auto count = 0;

            for (int j = 0; j < size; j++)
            {
                if (i == j) continue;
                if (std::abs(x - xs[j]) < avg_width && std::abs(y - ys[j]) < avg_height) count++;
            }

            if (count > (size / 4))
            {
                avg_x += x;
                avg_y += y;
                min_x = min_x < x ? min_x : x;
                min_y = min_y < y ? min_y : y;
                max_x = max_x < x ? x : max_x;
                max_y = max_y < y ? y : max_y;
                valid_size++;
            }
        }
        if (min_x > max_x || min_y > max_y) return nullptr;

        avg_x /= valid_size;
        avg_y /= valid_size;
        // std::cout << "min_x = " << min_x << ", min_y = " << min_y << ", max_x = " << max_x << ", max_y = " << max_y << ", avg_x = " << avg_x << ", avg_y = " << avg_y << ", size = " << valid_size << std::endl;

        int half_width  = (max_x - min_x) / 2;
        int half_height = (max_y - min_y) / 2;
        int step_x      = 0;
        int step_y      = 0;
        int rect[4]     = {0};
        do {
            if (rect[0] > min_x || rect[1] < max_x) step_x += std::max(1, (int)(5 * scale));
            if (rect[2] > min_y || rect[3] < max_y) step_y += std::max(1, (int)(5 * scale));
            rect[0] = std::max(0, avg_x - half_width - step_x);
            rect[1] = std::min(img_width, avg_x + half_width + step_x);
            rect[2] = std::max(0, avg_y - half_height - step_y);
            rect[3] = std::min(img_height, avg_y + half_height + step_y);
            // std::cout << "rect: " << rect[0] << " " << rect[1] << " " << rect[2] << " " << rect[3] << " " << step_x << " " << step_y << std::endl;
        } while (rect[0] > min_x || rect[1] < max_x || rect[2] > min_y || rect[3] < max_y);

        int x      = rect[0] / scale;
        int y      = rect[2] / scale;
        int width  = (rect[1] - rect[0]) / scale;
        int height = (rect[3] - rect[2]) / scale;

        static int s_results[4] = {0, 0, 0, 0};
        s_results[0]            = x;
        s_results[1]            = y;
        s_results[2]            = width;
        s_results[3]            = height;
        if (width <= 0 || height <= 0) return nullptr;
        // std::cout << "x: " << x << " y: " << y << " width: " << width << " height: " << height << std::endl;
        return s_results;
    }

    std::shared_ptr<DataImage> GetColorImage(std::shared_ptr<DataImage> img, int color_classify_index, float scale = 0.2)
    {
        auto img_width  = img->GetWidth();
        auto img_height = img->GetHeight();
        auto rects      = GetColorRect(img, color_classify_index, scale);
        if (rects == nullptr) return nullptr;
        auto x = std::max(0, rects[0] - 10);
        auto y = std::max(0, rects[1] - 10);
        auto w = std::min(img_width, rects[2] + 20);
        auto h = std::min(img_height, rects[3] + 20);
        return img->GetSubImage(x, y, w, h);
    }

    std::vector<std::string> GetClassifyList()
    {
        return std::move(ColorModel::GetClassifyList());
    }

    int *GetRedRect(std::shared_ptr<DataImage> img, float scale = 0.2)
    {
        img = img->Scale(scale, scale);

        std::vector<int> xs;
        std::vector<int> ys;

        auto img_width    = img->GetWidth();
        auto img_height   = img->GetHeight();
        auto img_channels = img->GetChannel();
        auto img_datas    = img->GetData();

        for (int y = 0; y < img_height; y++)
        {
            for (int x = 0; x < img_width; x++)
            {
                auto r     = img->GetPixel(x, y, 0) * 255;
                auto g     = img->GetPixel(x, y, 1) * 255;
                auto b     = img->GetPixel(x, y, 2) * 255;
                if ((r - g) > 0x20 && (r - b) > 0x20 && (g < b ? b - g : g - b) < 0x18)
                {
                    xs.push_back(x);
                    ys.push_back(y);
                }
            }
        }

        int size = xs.size();
        if (size == 0) return nullptr;

        auto avg_width  = 0;
        auto avg_height = 0;
        for (int i = 0; i < size; i++)
        {
            auto max_width  = 0;
            auto max_height = 0;
            for (int j = i + 1; j < size; j++)
            {
                max_width  = std::max(max_width, std::abs(xs[i] - xs[j]));
                max_height = std::max(max_height, std::abs(ys[i] - ys[j]));
            }
            avg_width += max_width;
            avg_height += max_height;
        }
        avg_width  = avg_width / size * 2;
        avg_height = avg_height / size * 2;

        auto avg_x      = 0;
        auto avg_y      = 0;
        auto min_x      = img_width;
        auto min_y      = img_height;
        auto max_x      = 0;
        auto max_y      = 0;
        auto valid_size = 0;
        for (int i = 0; i < size; i++)
        {
            auto x     = xs[i];
            auto y     = ys[i];
            auto count = 0;

            for (int j = 0; j < size; j++)
            {
                if (i == j) continue;
                if (std::abs(x - xs[j]) < avg_width && std::abs(y - ys[j]) < avg_height) count++;
            }

            if (count > (size / 4))
            {
                avg_x += x;
                avg_y += y;
                min_x = min_x < x ? min_x : x;
                min_y = min_y < y ? min_y : y;
                max_x = max_x < x ? x : max_x;
                max_y = max_y < y ? y : max_y;
                valid_size++;
            }
        }
        if (min_x > max_x || min_y > max_y) return nullptr;

        avg_x /= valid_size;
        avg_y /= valid_size;
        // std::cout << "min_x = " << min_x << ", min_y = " << min_y << ", max_x = " << max_x << ", max_y = " << max_y << ", avg_x = " << avg_x << ", avg_y = " << avg_y << ", size = " << valid_size << std::endl;

        int half_width  = (max_x - min_x) / 2;
        int half_height = (max_y - min_y) / 2;
        int step_x      = 0;
        int step_y      = 0;
        int rect[4]     = {0};
        do {
            if (rect[0] > min_x || rect[1] < max_x) step_x += std::max(1, (int)(5 * scale));
            if (rect[2] > min_y || rect[3] < max_y) step_y += std::max(1, (int)(5 * scale));
            rect[0] = std::max(0, avg_x - half_width - step_x);
            rect[1] = std::min(img_width, avg_x + half_width + step_x);
            rect[2] = std::max(0, avg_y - half_height - step_y);
            rect[3] = std::min(img_height, avg_y + half_height + step_y);
            // std::cout << "rect: " << rect[0] << " " << rect[1] << " " << rect[2] << " " << rect[3] << " " << step_x << " " << step_y << std::endl;
        } while (rect[0] > min_x || rect[1] < max_x || rect[2] > min_y || rect[3] < max_y);

        int x      = rect[0] / scale;
        int y      = rect[2] / scale;
        int width  = (rect[1] - rect[0]) / scale;
        int height = (rect[3] - rect[2]) / scale;

        static int s_results[4] = {0, 0, 0, 0};
        s_results[0]            = x;
        s_results[1]            = y;
        s_results[2]            = width;
        s_results[3]            = height;
        if (width <= 0 || height <= 0) return nullptr;
        // std::cout << "x: " << x << " y: " << y << " width: " << width << " height: " << height << std::endl;
        return s_results;
    }

};

#endif